import os
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('-i', '--img-dir', default='images', help='The directory of images')
parser.add_argument('-m', '--mask-dir', default='labels', help='The directory of mask')
parser.add_argument('-o', '--output', default='./test_list.txt', help='The path to save output')
args = parser.parse_args()

img_ext = '.tif'
mask_ext = '.png'
image_names = [name for name in os.listdir(args.img_dir) if name.endswith(img_ext)]
mask_names = [name for name in os.listdir(args.mask_dir) if name.endswith(mask_ext)]

common_names = [name[:-len(img_ext)] for name in image_names if name[:-len(img_ext)] + mask_ext in mask_names]
if len(common_names) != len(image_names) or len(common_names) != len(mask_names):
    print(f'WARNING: common: {len(common_names)}, images: {len(image_names)}, masks: {len(mask_names)}')
else:
    print(f'There are {len(common_names)} pictures with masks')

with open(args.output, 'w') as f:
    for name in common_names:
        f.write(f"{os.path.join(args.img_dir, name + img_ext)} {os.path.join(args.mask_dir, name + mask_ext)}\n")
